# -*- coding: utf-8 -*-

import rsa
from logger import Logger
from base64 import b64encode, b64decode

   
class GroupPublicKey:
    def __init__(self, *args, **kwargs):
        b64 = kwargs.get("b64")
        if b64 is not None:
            key = b64.split(" ")
            self.pk_e = rsa.PublicKey.load_pkcs1(b64decode(key[0]), format = "DER")
            self.pk_s = rsa.PublicKey.load_pkcs1(b64decode(key[1]), format = "DER")
        else:
            self.pk_e = args[0]
            self.pk_s = args[1]
    
    def __str__(self):
        key = b64encode(self.pk_e.save_pkcs1(format = "DER"))
        key += " "
        key += b64encode(self.pk_s.save_pkcs1(format = "DER"))
        return key

     
class GroupManagerSecretKey:
    def __init__(self, *args, **kwargs):
        b64 = kwargs.get("b64")
        if b64 is not None:
            key = b64.split(" ")
            self.sk_e = rsa.PrivateKey.load_pkcs1(b64decode(key.pop(0)), format = "DER")
            self.pk_i = {}
            for i in range(0, len(key), 2):
                self.pk_i[key[i]] = rsa.PublicKey.load_pkcs1(b64decode(key[i + 1]), format = "DER")
        else:
            self.sk_e = args[0]
            self.pk_i = {}
    
    def __str__(self):
        key = b64encode(self.sk_e.save_pkcs1(format = "DER"))
        for i in self.pk_i.keys():
            key += " " + i + " "
            key += b64encode(self.pk_i[i].save_pkcs1(format = "DER"))
        return key
        
    def add(self, ident, pk_i):
        self.pk_i[ident] = pk_i


class GroupSecretKey:
    def __init__(self, *args, **kwargs):
        b64 = kwargs.get("b64")
        if b64 is not None:
            key = b64.split(" ")
            self.ident = key[0]
            self.sk_i = rsa.PrivateKey.load_pkcs1(b64decode(key[1]), format = "DER")
            self.pk_e = rsa.PublicKey.load_pkcs1(b64decode(key[2]), format = "DER")
            self.sk_s = rsa.PrivateKey.load_pkcs1(b64decode(key[3]), format = "DER")
        else:
            self.ident = args[0]
            self.sk_i = args[1]
            self.pk_e = args[2]
            self.sk_s = args[3]
            
    def __str__(self):
        key = self.ident
        key += " " + b64encode(self.sk_i.save_pkcs1(format = "DER"))
        key += " " + b64encode(self.pk_e.save_pkcs1(format = "DER"))
        key += " " + b64encode(self.sk_s.save_pkcs1(format = "DER"))
        return key


def GKg(k, n):
    Logger.info("Generating group encryption keys... ")
    pk_e, sk_e = rsa.newkeys(2 * k)
    Logger.done()
    Logger.info("Generating group signing keys... ")
    pk_s, sk_s = rsa.newkeys(k)
    Logger.done()
    
    gpk = GroupPublicKey(pk_e, pk_s)
    gmsk = GroupManagerSecretKey(sk_e)
    gsk = {}
    
    for ident in n:
        Logger.info("Generating signing keys for user {}... ".format(ident))
        pk_i, sk_i = rsa.newkeys(k)
        gsk[ident] = GroupSecretKey(ident, sk_i, pk_e, sk_s)
        gmsk.add(ident, pk_i)
        Logger.done()

    return gpk, gmsk, gsk


def GSig(m, gsk):
    Logger.info("Creating user signature... ")
    s = b64encode(rsa.sign(m, gsk.sk_i, "SHA-256"))
    Logger.done()
    
    Logger.info("Encrypting user signature... ")
    C = b64encode(rsa.encrypt(gsk.ident + " " + s, gsk.pk_e))
    Logger.done()
    
    Logger.info("Creating group signature... ")
    gs = b64encode(rsa.sign(m + C, gsk.sk_s, "SHA-256"))
    Logger.done()
    
    sigma = gs + " " + C
    return sigma

def GVf(m, sigma, gpk):
    Logger.info("Verifying group signature... ")
    gs = b64decode(sigma.split(" ")[0])
    C = sigma.split(" ")[1]
    
    try:
        rsa.verify(m + C, gs, gpk.pk_s)
    except rsa.VerificationError:
        Logger.done()
        return False
    
    Logger.done()
    return True

def GOpen(m, sigma, gpk, gmsk):
    if not GVf(m, sigma, gpk):
        return None
    
    C = b64decode(sigma.split(" ")[1])
    i_s = rsa.decrypt(C, gmsk.sk_e).split(" ")
    ident = i_s[0]
    s = b64decode(i_s[1])
    
    Logger.info("Verifying user signature... ")
    try:
        rsa.verify(m, s, gmsk.pk_i[ident])
    except rsa.VerificationError:
        Logger.done()
        return None
    
    Logger.done()
    return ident

